import sys

from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt

import numpy as np

traj_filename = sys.argv[1]
output_flname = sys.argv[2]

traj_fl = open(traj_filename)

frames = {}

try:
    while True:
        num_objects, step = map(int, next(traj_fl).split())
        for j in xrange(num_objects):
            cols = next(traj_fl).split()
            name = cols[0]
            positions = np.array(map(float, cols[1:4]), np.float128)
            velocities = np.array(map(float, cols[4:7]), np.float128)
            if name not in frames:
                frames[name] = ([], [], [])
            frames[name][0].append(positions[0])
            frames[name][1].append(positions[1])
            frames[name][2].append(positions[2])
except StopIteration:
    pass

fig = plt.figure()
ax = axes3d.Axes3D(fig)
for name in ["Earth", "Mercury"]:
    x, y, z = frames[name]
    ax.plot3D(x,y,z)
plt.savefig(output_flname, DPI=200)


